from typing import Optional, Union
import random
import httpx

from axiom_boot.conf.manager import Settings
from axiom_boot.di import service, order, conditional_on_setting
from axiom_boot.logging.setup import get_logger
from axiom_boot.scraper.interfaces import ScraperMiddleware
from axiom_boot.scraper.models import Target, Response


logger = get_logger(__name__)


@order(50)
@service()
@conditional_on_setting(key="scraper__proxy_enabled", expected_value=True)
class ProxyMiddleware(ScraperMiddleware):
    _proxy_list: list[str] = []

    def __init__(self, settings: Settings):
        self._settings = settings.scraper
        if self._settings.proxy_mode == "fixed_list":
            self._proxy_list = self._settings.proxy_list
            if not self._proxy_list:
                raise ValueError("代理模式为 'fixed_list'，但 'proxy_list' 为空。")

    async def process_request(self, target: Target) -> Optional[Union[Target, Response]]:
        proxy = await self._get_proxy()
        if proxy:
            logger.debug(f"为 URL {target.url} 分配代理: {proxy}")
            new_metadata = target.metadata.copy()
            new_metadata["proxy"] = proxy
            return target.model_copy(update={"metadata": new_metadata})
        return target

    async def process_response(self, target: Target, response: Response) -> Optional[Union[Target, Response]]:
        return response

    async def process_exception(self, target: Target, exception: Exception) -> Optional[Union[Target, Response]]:
        return None

    async def _get_proxy(self) -> Optional[str]:
        if self._settings.proxy_mode == "fixed_list":
            return random.choice(self._proxy_list)
        elif self._settings.proxy_mode == "api":
            if not self._settings.proxy_api_url:
                logger.error("代理模式为 'api'，但 'proxy_api_url' 未配置。")
                return None
            try:
                async with httpx.AsyncClient() as client:
                    response = await client.get(self._settings.proxy_api_url)
                    response.raise_for_status()
                    return response.text.strip()
            except httpx.HTTPError as e:
                logger.error(f"从代理 API 获取代理失败: {e}")
                return None
        return None 